import os
import cv2
import math
import time
import rclpy
import threading
import numpy as np
import simpleaudio as sa

from loguru import logger
from rclpy.node import Node
from datetime import datetime
from collections import deque, Counter
from std_msgs.msg import String, Float64
from ais_bench.infer.interface import InferSession
from flask import Flask, Response
from pydub import AudioSegment


DEBUG = False
FRAME_WIDTH = 640
FRAME_HEIGHT = 480

app = Flask(__name__)
frame_queue = deque(maxlen=1)

audio_map = {
    "level": AudioSegment.from_wav("/home/Code/RoboticDog/resource/pre_audio_file/level2wav.wav"),
    "stairs": AudioSegment.from_wav("/home/Code/RoboticDog/resource/pre_audio_file/stairs2wav.wav"),
    "overleap": AudioSegment.from_wav("/home/Code/RoboticDog/resource/pre_audio_file/overleap2wav.wav"),
    "slope": AudioSegment.from_wav("/home/Code/RoboticDog/resource/pre_audio_file/slope2wav.wav"),
    "stop": AudioSegment.from_wav("/home/Code/RoboticDog/resource/pre_audio_file/stop2wav.wav"),
}


class Decider:
    def __init__(self, ):
        # 参数
        window_size = 10
        self.windows = [deque(maxlen=window_size), deque(maxlen=window_size)] # 图像分类窗口

        self.degree_queue = deque([0], maxlen=window_size)  # degree 缓存
        self.degree_range = [-30, -5, 5, 30] # 斜坡角度范围 [stop <- x0 -> slope <- x1 -> level <- x2 -> slop <- x3 -> stop]
        self.degree_info = 'None'
        
        # 置信度阈值
        self.threshold = [
            {
                0: 0.8,
                1: 0.8,
                2: 0.8,
            },
            {
                0: 0.65,
                1: 0.85,
                2: 0.85

            },
        ] 
        
        # 窗口中的稳定值
        self.stable_num = 7

        self.decision_map = {
            '00': 'level', # 两个都是平路
            '01': 'overleap',
            '02': 'overleap',
            '10': 'stairs',
            '11': 'stairs',
            '12': 'stairs',
            '20': 'stairs',
            '21': 'stairs',
            '22': 'stairs',
        }

        self.key = 'None'
        self.decision = 'None'
        self.status = 'level'

    def add_frame(self, pred_class, confidence):
        for i, pred in enumerate(pred_class):
            confd = confidence[i]
            # 筛选高置信度帧
            if confd >= self.threshold[i][pred]:
                self.windows[i].append(pred)

        self.update_status()

    def get_degree_info(self):
        mean_degree = sum(self.degree_queue) / len(self.degree_queue)

        # 临界值波动消除
        for edge in self.degree_range:
            if edge - 1 < mean_degree < edge + 1:
                return self.degree_info
        
        if mean_degree <= self.degree_range[0]:
            return 'stop'
        elif self.degree_range[0] < mean_degree <= self.degree_range[1]:
            return 'slope'
        elif self.degree_range[1] < mean_degree <= self.degree_range[2]:
            return 'level'
        elif self.degree_range[2] < mean_degree <= self.degree_range[3]:
            return 'slope'
        elif self.degree_range[3] < mean_degree:
            return 'stop'
        else:
            return 'None'
        
    def update_status(self):
        key = ['N', 'N']
        for i, window in enumerate(self.windows):
            if len(window) == 0:
                continue
            counts = Counter(window)
            stable_class, count = counts.most_common(1)[0]
            if count >= self.stable_num:
                key[i] = str(stable_class)
        key = ''.join(key)
        
        self.key = key if 'N' not in key else f'{key}:None'
        self.decision = self.decision_map[key] if 'N' not in key else self.decision # 判断平路、越障、楼梯模式
        self.degree_info = self.get_degree_info() # 判断平路、斜坡、停止模式
        
        if self.status == 'stairs':
            if self.degree_info == 'level': 
                self.status = self.decision
        elif self.status == 'overleap':
            if self.degree_info == 'stop':
                self.status = self.degree_info
            else:
                self.status = self.decision
        elif self.status == 'level':
            if self.decision == 'level':
                self.status = self.degree_info
            else:
                self.status = self.decision
        elif self.status == 'slope':
            if self.decision != 'level':
                self.status = self.decision
            else:
                self.status = self.degree_info
        elif self.status == 'stop':
            if self.degree_info == 'stop':
                self.status = self.degree_info
            else:
                if self.decision != 'level':
                    self.status = self.decision
                else:
                    self.status = self.degree_info
        else:
            self.status = self.degree_info


class RoadClass:
    def __init__(
        self, 
        om1_file='/home/Code/RoboticDog/resource/model/terrains1_bs1_448_aipp.om',
        om2_file='/home/Code/RoboticDog/resource/model/terrains2_bs1_448_aipp.om',
        save_img_path='/home/nzd/save_imgs',
        debug=False,
    ):
        
        # args
        device_id = 0
     
        self.box_x = int((FRAME_WIDTH - FRAME_HEIGHT) / 2)
        self.box_y = 0
        self.box_widht = FRAME_HEIGHT
        self.box_height = FRAME_HEIGHT
        self.resize = (448, 448)
                
        # init
        self.om_sess1 = InferSession(device_id, om1_file)
        self.om_sess2 = InferSession(device_id, om2_file)

        self.decider = Decider()

        # debug
        self.debug = debug
        self.save_img_path = save_img_path
        if not os.path.exists(self.save_img_path):
            os.makedirs(self.save_img_path)

    def preprocess(self, img):
        # 1. 提取框内图像
        cropped_img = img[self.box_y : self.box_y+self.box_height, self.box_x : self.box_x+self.box_widht]

        # 2. resize
        resized_img = cv2.resize(cropped_img, self.resize, interpolation=cv2.INTER_LINEAR)

        return resized_img

    def postpreocess(self, logits, axis=1):
        # 1. 数值稳定性处理：减去最大值防止指数溢出
        max_logits = np.max(logits, axis=axis, keepdims=True)
        exp_logits = np.exp(logits - max_logits)
        
        # 2. 计算 softmax 分母（沿指定轴求和）
        sum_exp = np.sum(exp_logits, axis=axis, keepdims=True)
        
        # 3. 计算 softmax 概率
        probabilities = exp_logits / sum_exp

        # 4. 获取分类indx，和概率
        label = np.argmax(probabilities, axis=1)[0]  # 每行最大值的索引
        prob = np.max(probabilities, axis=1)[0]      # 每行的最大值
         
        return label, prob

    def save_info(self, image, label1, prob1, label2, prob2, windows, key, decision, degree_queue, degree_info, status):
        mean_degree = sum(degree_queue) / len(degree_queue)

        debug_info = f"pred: {label1}:{prob1:.3f}_{label2}:{prob2:.3f}\nwindows: {windows}\nkey: {key}\ndecision: {decision}\nmean_degree: {mean_degree}\ndegree_info: {degree_info}\nstatus: {status}"
        print(debug_info)
        
        now = datetime.now()
        formatted_time = now.strftime("%Y_%m_%d_%H_%M_%S_%f")[:-3]
        cv2.imwrite(f"{self.save_img_path}/{formatted_time}_{label1}_{int(prob1*1000)}_{label2}_{int(prob2*1000)}_{int(mean_degree*100)}.jpg", image)

        return debug_info


    def infer(self, image):
        img = self.preprocess(image)
        
        pred1 = self.om_sess1.infer([img], custom_sizes=1000)[0]
        pred2 = self.om_sess2.infer([img], custom_sizes=1000)[0]

        label1, prob1 = self.postpreocess(pred1)
        label2, prob2 = self.postpreocess(pred2)
        
        self.decider.add_frame([label1, label2], [prob1, prob2])
                        
        if self.debug:
           debug_info = self.save_info(image, label1, prob1, label2, prob2, self.decider.windows, self.decider.key, self.decider.decision, self.decider.degree_queue, self.decider.degree_info, self.decider.status)
           return self.decider.status, debug_info

        return self.decider.status, None


class RoadClassification(Node):  # for inference
    def __init__(self, name, debug=False):
        super().__init__(name)
        self.get_logger().info("大家好, 我是%s节点！" % name)
        self.get_logger().info("初始化road_classification_node成功！")

        # 创建订阅者
        self.create_subscription(Float64, "pitch", self.listener_pitch_msg, 1)

        # 创建发布者, 发布数据
        self.cmd = String()
        self.road_result_pub = self.create_publisher(String, "cmd_switch_gait", 1)

        # 创建定时器
        self.create_timer(0.0001, self.infer_frame)
        
        # 创建摄像头
        self.frame = None
        self.camera = cv2.VideoCapture(0)  # 0通常是内置摄像头的索引
        if not self.camera.isOpened():
            raise RuntimeError("无法打开摄像头")
        self.camera.set(cv2.CAP_PROP_FRAME_WIDTH, FRAME_WIDTH)
        self.camera.set(cv2.CAP_PROP_FRAME_HEIGHT, FRAME_HEIGHT)
        
        # 创建实例
        self.debug = debug
        self.rc = RoadClass(debug=debug)
    
    def infer_frame(self, ):
        success, frame = self.camera.read()  # 从摄像头读取一帧

        if success:
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

            self.cmd.data, info = self.rc.infer(frame)

            logger.info(f"==========road_classification: {self.cmd.data}")
            if self.debug:
                frame_queue.append([frame, info])
                logger.info(f"==========road_classification: [INFO]{info}")
            else:
                self.road_result_pub.publish(self.cmd)
                pass
        else: 
            logger.info(f"==========road_classification: camera.read failse")
        

    def listener_pitch_msg(self, msg):
        degree = msg.data * (180 / math.pi)
        logger.info(f'====================dog degree: {degree}')
        self.rc.decider.degree_queue.append(degree)


def gen_frames():
    box_x, box_y = int((FRAME_WIDTH - FRAME_HEIGHT) / 2), 0
    box_width, box_height = FRAME_HEIGHT, FRAME_HEIGHT
    start_y = 20
    line_h = 18

    status = "None"

    while True:
        if len(frame_queue) != 0:
            frame, info = frame_queue[0]
            cv2.rectangle(frame, (box_x, box_y), 
                (box_x + box_width, box_y + box_height), 
                (0, 255, 0), 2)
            
            info = info.split('\n')
            results = info[-1].split(': ')[-1]
            infos = info[:4]
            window1 = f"window1: {infos[1].split(', maxlen')[0].replace('windows: [deque(', '')}"
            window2 = f"window2: {infos[1].split('deque(')[-1].replace(', maxlen=10)]', '')}"
            infos_ = [infos[0], window1, window2, infos[2], infos[3], info[4], info[5], info[6]]
           
            for i, line in enumerate(infos_):
                y = start_y + i*line_h
                cv2.putText(frame, line, (10, y), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1, cv2.LINE_AA)
            
            ret, buffer = cv2.imencode('.jpg', frame)
            frame = buffer.tobytes()

            if status != results:
                status = results
                if status != 'None':
                    play_obj = sa.play_buffer(audio_map[status].raw_data, num_channels=audio_map[status].channels, bytes_per_sample=audio_map[status].sample_width, sample_rate=audio_map[status].frame_rate)
                    
            yield (b'--frame\r\n'
                   b'Content-Type: image/jpeg\r\n\r\n' + frame + b'\r\n')  # 发送帧数据


@app.route('/video_feed')
def video_feed():
    return Response(gen_frames(), mimetype='multipart/x-mixed-replace; boundary=frame')


def run_ros_node():
    rclpy.init(args=None)
    road_classification_node = RoadClassification("road_classification_node", debug=DEBUG)
    rclpy.spin(road_classification_node)
    rclpy.shutdown()


def run_flask():
    print("[INFO]Access: http://192.168.1.100:5000/video_feed")
    app.run(host='192.168.1.100', port=5000)  # 在所有IP上运行，端口5000


if __name__ == '__main__':
    if DEBUG:
        # 创建并启动ROS线程
        ros_thread = threading.Thread(target=run_ros_node)
        ros_thread.daemon = True  # 设为守护线程，主退出时自动终止
        ros_thread.start()

        # 在主线程运行Flask
        run_flask()
    else:
        run_ros_node()
